import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

df = pd.read_csv('./files/通威股份.csv')[['open','close','low','high','date']]
df['date'] = pd.to_datetime(df['date'])
df['ma5'] = np.nan
df['ma30'] = np.nan

df['dates'] = df['date']
df.set_index('dates', inplace=True)

# 设定你想要筛选的日期范围
# start_date = '2021-01-02'
# end_date = '2017-12-15'
end_date = '2024-7-18'
 
# 使用布尔索引进行筛选
# mask = (df['date'] >= start_date) & (df['date'] <= end_date)
mask = df['date'] <= end_date
df = df.loc[mask]
# print(df)

for i in range(4, len(df)):
  df.loc[df.index[i], 'ma5'] = df['close'][i-4:i+1].mean()

for i in range(29, len(df)):
  df.loc[df.index[i], 'ma30'] = df['close'][i-29:i+1].mean()

# 简单方法求均线
df['ma5_T'] = df['close'].rolling(window=5).mean()
df['ma30_T'] = df['close'].rolling(window=30).mean()

print(df)
df = df[-100:]
df[['close','ma5','ma30']].plot()
plt.show()
